"""
Neo4j 数据库客户端
用于管理数据血缘图的存储和查询
"""

from typing import Dict, List, Any, Optional
import json
import hashlib
from datetime import datetime
from neo4j import GraphDatabase, Driver, Session
from neo4j.exceptions import ServiceUnavailable, TransientError

from app.core.config import settings
from app.core.logger import logger


class Neo4jClient:
    """Neo4j 客户端类"""
    
    def __init__(self):
        self.driver: Optional[Driver] = None
        self._connect()
    
    def _connect(self):
        """连接到Neo4j数据库"""
        try:
            self.driver = GraphDatabase.driver(
                settings.NEO4J_URI,
                auth=(settings.NEO4J_USER, settings.NEO4J_PASSWORD),
                max_connection_lifetime=3600,
                max_connection_pool_size=50,
                connection_acquisition_timeout=60
            )
            # 测试连接
            with self.driver.session() as session:
                session.run("RETURN 1")
            logger.info("✅ Neo4j连接成功")
        except Exception as e:
            logger.error(f"❌ Neo4j连接失败: {str(e)}")
            self.driver = None
    
    def close(self):
        """关闭连接"""
        if self.driver:
            self.driver.close()
            logger.info("Neo4j连接已关闭")
    
    def is_connected(self) -> bool:
        """检查连接状态"""
        if not self.driver:
            return False
        try:
            with self.driver.session() as session:
                session.run("RETURN 1")
            return True
        except:
            return False
    
    def execute_query(self, query: str, parameters: Dict = None) -> List[Dict]:
        """执行查询"""
        if not self.driver:
            raise Exception("Neo4j未连接")
        
        try:
            with self.driver.session() as session:
                result = session.run(query, parameters or {})
                return [record.data() for record in result]
        except Exception as e:
            logger.error(f"Neo4j查询执行失败: {str(e)}")
            raise
    
    def execute_write_transaction(self, queries: List[Dict]) -> Dict:
        """执行写事务"""
        if not self.driver:
            raise Exception("Neo4j未连接")
        
        def _tx_function(tx):
            results = []
            for query_info in queries:
                query = query_info.get('query')
                parameters = query_info.get('parameters', {})
                result = tx.run(query, parameters)
                results.append({
                    'query': query,
                    'summary': result.consume().counters
                })
            return results
        
        try:
            with self.driver.session() as session:
                results = session.execute_write(_tx_function)
                return {
                    'success': True,
                    'results': results
                }
        except Exception as e:
            logger.error(f"Neo4j写事务执行失败: {str(e)}")
            return {
                'success': False,
                'error': str(e)
            }
    
    def create_table_node(self, table_info: Dict) -> Dict:
        """创建表节点"""
        query = """
        MERGE (t:Table {id: $id})
        SET t.name = $name,
            t.database = $database,
            t.schema = $schema,
            t.table_type = $table_type,
            t.connector_type = $connector_type,
            t.connector_config = $connector_config,
            t.description = $description,
            t.updated_at = datetime()
        ON CREATE SET t.created_at = datetime()
        RETURN t
        """
        
        parameters = {
            'id': table_info['id'],
            'name': table_info['name'],
            'database': table_info.get('database', ''),
            'schema': table_info.get('schema', ''),
            'table_type': table_info.get('table_type', 'TABLE'),
            'connector_type': table_info.get('connector_type', ''),
            'connector_config': table_info.get('connector_config', {}),
            'description': table_info.get('description', '')
        }
        
        return self.execute_query(query, parameters)
    
    def create_column_node(self, column_info: Dict) -> Dict:
        """创建字段节点"""
        query = """
        MERGE (c:Column {id: $id})
        SET c.name = $name,
            c.data_type = $data_type,
            c.is_primary_key = $is_primary_key,
            c.is_nullable = $is_nullable,
            c.default_value = $default_value,
            c.comment = $comment,
            c.updated_at = datetime()
        ON CREATE SET c.created_at = datetime()
        RETURN c
        """
        
        parameters = {
            'id': column_info['id'],
            'name': column_info['name'],
            'data_type': column_info.get('data_type', ''),
            'is_primary_key': column_info.get('is_primary_key', False),
            'is_nullable': column_info.get('is_nullable', True),
            'default_value': column_info.get('default_value', ''),
            'comment': column_info.get('comment', '')
        }
        
        return self.execute_query(query, parameters)
    
    def create_task_node(self, task_info: Dict) -> Dict:
        """创建任务节点"""
        query = """
        MERGE (t:Task {id: $id})
        SET t.name = $name,
            t.sql_content = $sql_content,
            t.task_type = $task_type,
            t.updated_at = datetime()
        ON CREATE SET t.created_at = datetime()
        RETURN t
        """
        
        parameters = {
            'id': task_info['id'],
            'name': task_info['name'],
            'sql_content': task_info.get('sql_content', ''),
            'task_type': task_info.get('task_type', 'FLINK_SQL')
        }
        
        return self.execute_query(query, parameters)
    
    def create_lineage_relationship(self, rel_info: Dict) -> Dict:
        """创建血缘关系"""
        rel_type = rel_info['relationship_type']
        
        if rel_type == 'FLOWS_TO':
            query = """
            MATCH (source {id: $source_id}), (target {id: $target_id})
            MERGE (source)-[r:FLOWS_TO {task_id: $task_id}]->(target)
            SET r.sql_fragment = $sql_fragment,
                r.operation_type = $operation_type,
                r.updated_at = datetime()
            ON CREATE SET r.created_at = datetime()
            RETURN r
            """
        elif rel_type == 'MAPS_TO':
            query = """
            MATCH (source:Column {id: $source_id}), (target:Column {id: $target_id})
            MERGE (source)-[r:MAPS_TO {task_id: $task_id}]->(target)
            SET r.transformation = $transformation,
                r.operation_type = $operation_type,
                r.updated_at = datetime()
            ON CREATE SET r.created_at = datetime()
            RETURN r
            """
        elif rel_type == 'CONTAINS':
            query = """
            MATCH (table:Table {id: $source_id}), (column:Column {id: $target_id})
            MERGE (table)-[r:CONTAINS]->(column)
            SET r.updated_at = datetime()
            ON CREATE SET r.created_at = datetime()
            RETURN r
            """
        elif rel_type == 'DEPENDS_ON':
            query = """
            MATCH (task:Task {id: $source_id}), (table:Table {id: $target_id})
            MERGE (task)-[r:DEPENDS_ON {access_type: $access_type}]->(table)
            SET r.updated_at = datetime()
            ON CREATE SET r.created_at = datetime()
            RETURN r
            """
        else:
            raise ValueError(f"不支持的关系类型: {rel_type}")
        
        parameters = {
            'source_id': rel_info['source_id'],
            'target_id': rel_info['target_id'],
            'task_id': rel_info.get('task_id', ''),
            'sql_fragment': rel_info.get('sql_fragment', ''),
            'operation_type': rel_info.get('operation_type', ''),
            'transformation': rel_info.get('transformation', ''),
            'access_type': rel_info.get('access_type', '')
        }
        
        return self.execute_query(query, parameters)
    
    def get_lineage_upstream(self, node_id: str, depth: int = 3) -> List[Dict]:
        """获取上游血缘"""
        query = """
        MATCH path = (target {id: $node_id})<-[:FLOWS_TO*1..$depth]-(source)
        RETURN path, length(path) as depth
        ORDER BY depth
        """
        
        return self.execute_query(query, {'node_id': node_id, 'depth': depth})
    
    def get_lineage_downstream(self, node_id: str, depth: int = 3) -> List[Dict]:
        """获取下游血缘"""
        query = """
        MATCH path = (source {id: $node_id})-[:FLOWS_TO*1..$depth]->(target)
        RETURN path, length(path) as depth
        ORDER BY depth
        """
        
        return self.execute_query(query, {'node_id': node_id, 'depth': depth})
    
    def get_full_lineage_graph(self, limit: int = 1000) -> Dict:
        """获取完整血缘图"""
        # 获取所有节点
        nodes_query = """
        MATCH (n)
        RETURN n, labels(n) as labels
        LIMIT $limit
        """
        
        # 获取所有关系
        edges_query = """
        MATCH (source)-[r]->(target)
        RETURN source.id as source_id, target.id as target_id, 
               type(r) as relationship_type, properties(r) as properties
        LIMIT $limit
        """
        
        nodes = self.execute_query(nodes_query, {'limit': limit})
        edges = self.execute_query(edges_query, {'limit': limit})
        
        return {
            'nodes': nodes,
            'edges': edges,
            'stats': {
                'node_count': len(nodes),
                'edge_count': len(edges)
            }
        }
    
    def clear_all_data(self) -> Dict:
        """清空所有数据（谨慎使用）"""
        query = "MATCH (n) DETACH DELETE n"
        return self.execute_write_transaction([{'query': query}])


# 全局Neo4j客户端实例
neo4j_client = Neo4jClient()
